package com.ruleengine.worker.runner.chain;

import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_RUNNER_DISABLED;
import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_RUNNER_DISPATCH_FAILED;
import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_RUNNER_NOT_ALLOWED_CLOSE;
import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_RUNNER_START_FAILED;
import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_RUNNER_STOP_FAILED;
import static com.ruleengine.worker.model.Status.Phase.BANNED;
import static com.ruleengine.worker.model.Status.Phase.DELETED;
import static com.ruleengine.worker.model.Status.Phase.ERROR;
import static com.ruleengine.worker.model.Status.Phase.PENDING;
import static com.ruleengine.worker.model.Status.Phase.RUNNING;
import static com.ruleengine.worker.runner.chain.ChainRunnerUtil.closeLinkedRunners;
import static com.ruleengine.worker.runner.chain.ChainRunnerUtil.getLinkedPubSubRunners;
import static com.ruleengine.worker.runner.chain.ChainRunnerUtil.initChainRunner;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.atomic.AtomicReference;

import com.lmax.disruptor.BusySpinWaitStrategy;
import com.lmax.disruptor.EventFactory;
import com.lmax.disruptor.EventHandler;
import com.lmax.disruptor.ExceptionHandler;
import com.lmax.disruptor.InsufficientCapacityException;
import com.lmax.disruptor.RingBuffer;
import com.lmax.disruptor.dsl.Disruptor;
import com.lmax.disruptor.dsl.ProducerType;
import com.ruleengine.destination.DestinationManager;
import com.ruleengine.worker.exception.ChainRunnerException;
import com.ruleengine.worker.model.Chain;
import com.ruleengine.worker.model.Status;
import com.ruleengine.worker.runner.rule.PubSubRuleRunner;
import com.ruleengine.worker.runner.rule.RuleRunner;
import com.ruleengine.worker.runner.rule.RuleRunnerFactory;

import cn.hutool.core.thread.ThreadUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

/**
 * @author lianghanwei on 2024/3/3
 */
@Slf4j
public class DisruptorChainRunner<T> implements ChainRunner<T> {

    /**
     * bufferSize must be a power of 2
     */
    private static final int MAX_BUFFER_CAPACITY = 1024;
    private static final String EXECUTOR_NAME_FORMAT = "chain-runner-disruptor-executor-%d";

    /**
     *
     */
    private final String chainName;

    private final Chain chain;

    private final AtomicReference<Status.Phase> statusPhase = new AtomicReference<>(PENDING);

    // TODO-lhw-2024/03/03: TBD
    private final ChainRunnerObserver chainRunnerObserver;

    /**
     *
     */
//    private final RuleRunnerFactory ruleRunnerFactory;

    private final PubSubRuleRunner<T, T> dispatchRunner;

    private final Map<String, RuleRunner> ruleUidRunnerMap;

    /**
     *
     */
    private final Disruptor<ItemEvent<T>> disruptor;

    private final RingBuffer<ItemEvent<T>> ringBuffer;

    private final EventFactory<ItemEvent<T>> itemEventFactory = ItemEvent::new;

    private final EventHandler<ItemEvent<T>> itemEventHandler =
        (itemEvent, sequence, endOfBatch) -> dispatch(itemEvent.getItem());

    private final ExceptionHandler<ItemEvent<T>> itemEventExceptionHandler = new ExceptionHandler<>() {
        @Override
        public void handleEventException(Throwable ex, long sequence, ItemEvent<T> event) {
            log.error("{}, name={}, sequence={}", CHAIN_RUNNER_DISPATCH_FAILED, chainName, sequence, ex);
        }

        @Override
        public void handleOnStartException(Throwable ex) {
            log.error("{}, name={}", CHAIN_RUNNER_START_FAILED, chainName, ex);
            throw new ChainRunnerException(CHAIN_RUNNER_START_FAILED);
        }

        @Override
        public void handleOnShutdownException(Throwable ex) {
            log.error("{}, name={}", CHAIN_RUNNER_STOP_FAILED, chainName, ex);
            throw new ChainRunnerException(CHAIN_RUNNER_STOP_FAILED);
        }
    };

    public DisruptorChainRunner(Chain chain, DestinationManager destinationManager) {
        this(chain, null, destinationManager);
    }

    public DisruptorChainRunner(Chain chain,
                                ChainRunnerObserver chainRunnerObserver,
                                DestinationManager destinationManager) {
        this.disruptor = new Disruptor<>(this.itemEventFactory, MAX_BUFFER_CAPACITY,
            ThreadUtil.newNamedThreadFactory(EXECUTOR_NAME_FORMAT, false),
            ProducerType.MULTI, new BusySpinWaitStrategy());
        this.disruptor.handleEventsWith(itemEventHandler);
        this.disruptor.handleExceptionsFor(itemEventHandler).with(itemEventExceptionHandler);

        this.ringBuffer = disruptor.getRingBuffer();

        this.chainName = chain.getUniqueName();
        this.chain = chain;
        this.chainRunnerObserver = chainRunnerObserver;

//        this.ruleRunnerFactory = ruleRunnerFactory;
        this.dispatchRunner = (PubSubRuleRunner) RuleRunnerFactory
                .newInstance(ChainRunnerUtil.createDispatchRule(chainName), destinationManager);
        this.ruleUidRunnerMap = new HashMap<>();

        initChainRunner(ruleUidRunnerMap, chain, dispatchRunner, destinationManager);
//        notifyObserver(chainRunnerObserver, this);
    }

    @Override
    public void start() {
        if (statusPhase.compareAndSet(PENDING, RUNNING)) {
            disruptor.start();
            log.info("Successfully start the runner, name={}", chainName);
//            notifyObserver(chainRunnerObserver, this);
        }
    }

    @Override
    public void stop() {
        if (statusPhase.compareAndSet(RUNNING, PENDING)) {
            log.info("Successfully stop the runner, name={}", chainName);
            // disable notify observer, the stop chain status is controlled by Flow Admin
        }
    }

    @Override
    public void close() {
        log.info("Start to close chain runner, name={}", chainName);

        if (statusPhase.get() == DELETED) {
            return;
        }

        if (statusPhase.compareAndSet(PENDING, DELETED)
            || statusPhase.compareAndSet(ERROR, DELETED)
            || statusPhase.compareAndSet(BANNED, DELETED)) {
            log.info("Successfully set the runner status as DELETED, name={}, status={}", chainName, getStatus());
        } else {
            throw new ChainRunnerException(CHAIN_RUNNER_NOT_ALLOWED_CLOSE);
        }

        // Waits until all ItemEvents currently in the disruptor have been processed
        disruptor.shutdown();
        // TODO-lhw-2024/03/04: 分不同用途的单例线程池
        closeLinkedRunners(chainName, getLinkedPubSubRunners(ruleUidRunnerMap), dispatchRunner, ForkJoinPool.commonPool()).join();


        log.info("Close chain runner successfully, name={}", chainName);
    }

    @Override
    public Status getStatus() {
        return Status.builder()
            .phase(statusPhase.get())
            .build();
    }

    @Override
    public void process(T item) {
        process(item, false);
    }

    public void process(T item, boolean isAllowedDrop) {
        if (statusPhase.get() == RUNNING || statusPhase.get() == PENDING) {
            publishItem(item, isAllowedDrop);
        } else {
            throw new ChainRunnerException(CHAIN_RUNNER_DISABLED);
        }
    }

    @Override
    public String getChainName() {
        return chainName;
    }

    @Override
    public String getUid() {
        return chain.getMetadata().getUid();
    }

    @Override
    public Chain getChain() {
        return chain;
    }

    @Override
    public String getVersion() {
        return chain.getMetadata().getResourceVersion();
    }

    @Override
    public int getBufferSize() {
        return (int) (ringBuffer.getBufferSize() - ringBuffer.remainingCapacity());
    }

    private void publishItem(T item, boolean isAllowedDrop) {
        long availableSequence;

        if (isAllowedDrop) {
            try {
                availableSequence = ringBuffer.tryNext();
            } catch (InsufficientCapacityException e) {
                log.warn("Chain runner drop item due to buffer full, current buffer size={}", getBufferSize());
                return;
            }
        } else {
            // block when ring buffer full
            availableSequence = ringBuffer.next();
        }

        try {
            ItemEvent<T> itemEvent = ringBuffer.get(availableSequence);
            itemEvent.setItem(item);
        } finally {
            ringBuffer.publish(availableSequence);
        }
    }

    private void dispatch(T item) {
        dispatchRunner.write(item)
            .onFailure(ex -> log.error(CHAIN_RUNNER_DISPATCH_FAILED, ex));

//        if (this.chainRunnerObserver != null) {
//            this.chainRunnerObserver.onSubmitComplete(chain, getRingBufferUsedSize() + lag);
//        }
    }

    @Data
    private static class ItemEvent<T> {
        private T item;
    }
}